#!/bin/bash

# Foldseek数据库自动构建脚本
# 功能：扫描PDB文件并构建Foldseek搜索数据库

set -e

# ============ 配置参数 ============
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
PDB_DIR="$PROJECT_ROOT/../data/pdb"
FOLDSEEK_DB_DIR="$PROJECT_ROOT/../foldseek_db"
LOG_FILE="$FOLDSEEK_DB_DIR/build_foldseek.log"

# ============ 函数定义 ============
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a "$LOG_FILE"
}

error_exit() {
    log "ERROR: $1"
    exit 1
}

# ============ 主流程 ============
log "========== 开始构建Foldseek数据库 =========="

# 1. 创建目录
log "创建Foldseek数据库目录..."
mkdir -p "$FOLDSEEK_DB_DIR"
cd "$FOLDSEEK_DB_DIR"

# 2. 检查Foldseek是否安装
if ! command -v foldseek &> /dev/null; then
    error_exit "Foldseek未安装，请先安装"
fi

log "Foldseek版本: $(foldseek --version)"

# 3. 检查PDB文件
log "检查PDB文件..."
if [ ! -d "$PDB_DIR" ]; then
    error_exit "PDB目录不存在: $PDB_DIR"
fi

PDB_COUNT=$(find "$PDB_DIR" -name "*.pdb" | wc -l)
log "找到 $PDB_COUNT 个PDB文件"

if [ "$PDB_COUNT" -eq 0 ]; then
    error_exit "未找到PDB文件"
fi

# 4. 创建PDB文件列表
log "创建PDB文件列表..."
find "$PDB_DIR" -name "*.pdb" > pdb_list.txt
log "PDB文件列表已创建"

# 5. 构建Foldseek数据库
log "构建Foldseek数据库（这可能需要一些时间）..."
foldseek createdb "$PDB_DIR" targetDB \
    --threads $(nproc) \
    2>&1 | tee createdb.log

if [ ${PIPESTATUS[0]} -ne 0 ]; then
    error_exit "Foldseek数据库构建失败，查看 createdb.log"
fi

# 6. 创建索引（加速搜索）
log "创建搜索索引..."
mkdir -p tmp
foldseek createindex targetDB tmp \
    --threads $(nproc) \
    2>&1 | tee createindex.log

if [ ${PIPESTATUS[0]} -ne 0 ]; then
    log "WARNING: 索引创建失败，但数据库仍可使用"
fi

# 7. 清理临时文件
log "清理临时文件..."
rm -rf tmp

# 8. 验证数据库
log "验证Foldseek数据库..."
DB_SIZE=$(du -sh targetDB | awk '{print $1}')
log "数据库大小: $DB_SIZE"

# 9. 生成元数据
log "生成元数据..."
cat > metadata.json <<EOF
{
  "build_date": "$(date -Iseconds)",
  "pdb_file_count": $PDB_COUNT,
  "database_type": "Foldseek",
  "database_size": "$DB_SIZE",
  "pdb_directory": "$PDB_DIR",
  "foldseek_version": "$(foldseek --version)"
}
EOF

log "========== Foldseek数据库构建完成 =========="
log "数据库位置: $FOLDSEEK_DB_DIR/targetDB"
log "日志文件: $LOG_FILE"

